"""
Use the post-rectified graphs to extract lines and measure orgonality between them.
"""
import torch
import numpy as np
from .calculate_base import H_solver, cross_mat
import os


def H_frobenius(H: np.ndarray, H_esti: np.ndarray) -> float:
    return np.sqrt(np.sum(H - H_esti) ** 2)


def H_hilbert_schmidt(H: np.ndarray, H_esti: np.ndarray) -> float:
    """
    Using Hilbert-Schmidt norm to measure the difference between two H matrix.
    """
    E = H - H_esti
    [U, D, Vh] = np.linalg.svd(E)
    # the torlerance of eigen value equal to zero is 1e-5
    ind = D > 1e-5
    Av = U[:, ind]
    u = Vh[~ind, :].T
    hs = 0
    for i in range(Av.shape[1]):
        for j in range(u.shape[1]):
            hs += np.sum((Av[:, i] - u[:, j]) ** 2)
    return np.sqrt(hs)


def HSIC(H: np.ndarray, H_esti: np.ndarray) -> float:
    D = np.identity(3) - 1 / 3 * np.ones((3, 3))
    return 1 / 4 * np.trace(H @ H.T @ D @ H_esti @ H_esti.T @ D)


def OFC(H_: np.ndarray, lines: list) -> float:
    lines = np.array(lines)
    line1 = lines[:, 0:3] @ H_.T
    line2 = lines[:, 3:6] @ H_.T
    line1_xy = (
        line1[:, 0:2] / np.sqrt(np.sum(line1[:, 0:2] ** 2, axis=1))[:, np.newaxis]
    )
    line2_xy = (
        line2[:, 0:2] / np.sqrt(np.sum(line2[:, 0:2] ** 2, axis=1))[:, np.newaxis]
    )
    error = np.sum(line1_xy * line2_xy, axis=1)
    return np.sqrt(np.mean(error**2))


def evaluate_main(
    lines: torch.TensorType,
    src_img: torch.TensorType,
    line_id: list,
    H: np.ndarray,
    lines_proj: torch.TensorType,
    output_dir: str = "./output/",
    img_name: str = "rectified_img",
) -> float:
    """
    Arguments:
      - `lines` : a tensor of shape (n, 6)，是校正图像上的直线，用于度量相似性
      - `lines_proj`: a tensor of shape (n, 6)，是原图上的直线，用于校正图像
    """
    # 直接使用现成的直线解出的结果不可靠，最好还是用原图再解一遍
    H_ = H_solver(
        cross_mat(torch.tensor(lines_proj, dtype=torch.float32, device="cuda"))
        .detach()
        .clone()
    )
    H_esti = torch.inverse(H_).cpu().numpy()
    # calculate the similarity of H matrix
    H_frobenius_norm = H_frobenius(H, H_esti)
    H_hs_norm = H_hilbert_schmidt(H, H_esti)
    H_HSIC = HSIC(H, H_esti)
    H_OFC = OFC(H_.cpu().numpy(), lines)
    print(f"Frobenius norm|Hilbert-Schmidt norm|OFC|HSIC")
    output_str = f"{H_frobenius_norm:.5f}\t{H_hs_norm:.5f}\t{H_OFC:.5f}\t{H_HSIC:.5f}"
    print(output_str)
    # copy the output str to clipboard conveniently
    os.system(f"echo {output_str} | clip")


if __name__ == "__main__":
    H = np.array(
        [
            [8.00e-01, 2.00e-01, 0.00e00],
            [0.00e00, 1.25e00, 0.00e00],
            [1.00e-04, 1.00e-04, 1.00e00],
        ]
    )
    H_esti = np.array(
        [
            [8.0214e-01, 2.0923e-01, 0.0000e00],
            [0.0000e00, 1.2467e00, 0.0000e00],
            [4.3511e-05, 1.8629e-04, 1.0000e00],
        ]
    )
    print(H_frobenius(H, H_esti))
    print(H_hilbert_schmidt(H, H_esti))
    print(HSIC(H, H_esti))
    print(OFC(H, H_esti))
